#
# SPDX-FileCopyrightText: Copyright (c) 2021-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0#
"""TB configuration for the NR (5G) module of Sionna PHY"""

from .config import Config
from .utils import decode_mcs_index

class TBConfig(Config):
    # pylint: disable=line-too-long
    r"""
    The TBConfig objects sets parameters related to the transport block
    encoding, as described in TS 38.214 [3GPP38214]_.

    All configurable properties can be provided as keyword arguments during the
    initialization or changed later.

    The TBConfig is configured by selecting the modulation and coding scheme
    (MCS) tables and index.

    Example
    -------
    >>> tb_config = TBConfig(mcs_index=13)
    >>> tb_config.mcs_table = 3
    >>> tb_config.channel_type = "PUSCH"
    >>> tb_config.show()

    The following tables provide an overview of the corresponding coderates and
    modulation orders.

    .. table:: MCS Index Table 1 (Table 5.1.3.1-1 in [3GPP38214]_)
        :align: center

        +-------------------+--------------------+-------------------------+-----------------------+
        | | MCS Index       | | Modulation Order | | Target Coderate       | | Spectral Efficiency |
        | | :math:`I_{MCS}` | | :math:`Q_m`      | | :math:`R\times[1024]` | |                     |
        +===================+====================+=========================+=======================+
        | 0                 | 2                  | 120                     | 0.2344                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 1                 | 2                  | 157                     | 0.3066                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 2                 | 2                  | 193                     | 0.3770                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 3                 | 2                  | 251                     | 0.4902                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 4                 | 2                  | 308                     | 0.6016                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 5                 | 2                  | 379                     | 0.7402                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 6                 | 2                  | 449                     | 0.8770                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 7                 | 2                  | 526                     | 1.0273                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 8                 | 2                  | 602                     | 1.1758                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 9                 | 2                  | 679                     | 1.3262                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 10                | 4                  | 340                     | 1.3281                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 11                | 4                  | 378                     | 1.4766                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 12                | 4                  | 434                     | 1.6953                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 13                | 4                  | 490                     | 1.9141                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 14                | 4                  | 553                     | 2.1602                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 15                | 4                  | 616                     | 2.4063                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 16                | 4                  | 658                     | 2.5703                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 17                | 6                  | 438                     | 2.5664                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 18                | 6                  | 466                     | 2.7305                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 19                | 6                  | 517                     | 3.0293                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 20                | 6                  | 567                     | 3.3223                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 21                | 6                  | 616                     | 3.6094                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 22                | 6                  | 666                     | 3.9023                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 23                | 6                  | 719                     | 4.2129                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 24                | 6                  | 772                     | 4.5234                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 25                | 6                  | 822                     | 4.8164                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 26                | 6                  | 873                     | 5.1152                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 27                | 6                  | 910                     | 5.3320                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 28                | 6                  | 948                     | 5.5547                |
        +-------------------+--------------------+-------------------------+-----------------------+

    .. table:: MCS Index Table 2 (Table 5.1.3.1-2 in [3GPP38214]_)
        :align: center

        +-------------------+--------------------+-------------------------+-----------------------+
        | | MCS Index       | | Modulation Order | | Target Coderate       | | Spectral Efficiency |
        | | :math:`I_{MCS}` | | :math:`Q_m`      | | :math:`R\times[1024]` | |                     |
        +===================+====================+=========================+=======================+
        | 0                 | 2                  | 120                     | 0.2344                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 1                 | 2                  | 193                     | 0.3770                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 2                 | 2                  | 308                     | 0.6016                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 3                 | 2                  | 449                     | 0.8770                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 4                 | 2                  | 602                     | 1.1758                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 5                 | 4                  | 378                     | 1.4766                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 6                 | 4                  | 434                     | 1.6953                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 7                 | 4                  | 490                     | 1.9141                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 8                 | 4                  | 553                     | 2.1602                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 9                 | 4                  | 616                     | 2.4063                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 10                | 4                  | 658                     | 2.5703                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 11                | 6                  | 466                     | 2.7305                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 12                | 6                  | 517                     | 3.0293                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 13                | 6                  | 567                     | 3.3223                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 14                | 6                  | 616                     | 3.6094                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 15                | 6                  | 666                     | 3.9023                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 16                | 6                  | 719                     | 4.2129                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 17                | 6                  | 772                     | 4.5234                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 18                | 6                  | 822                     | 4.8164                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 19                | 6                  | 873                     | 5.1152                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 20                | 8                  | 682.5                   | 5.3320                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 21                | 8                  | 711                     | 5.5547                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 22                | 8                  | 754                     | 5.8906                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 23                | 8                  | 797                     | 6.2266                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 24                | 8                  | 841                     | 6.5703                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 25                | 8                  | 885                     | 6.9141                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 26                | 8                  | 916.5                   | 7.1602                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 27                | 8                  | 948                     | 7.4063                |
        +-------------------+--------------------+-------------------------+-----------------------+

    .. table:: MCS Index Table 3 (Table 5.1.3.1-3 in [3GPP38214]_)
        :align: center

        +-------------------+--------------------+-------------------------+-----------------------+
        | | MCS Index       | | Modulation Order | | Target Coderate       | | Spectral Efficiency |
        | | :math:`I_{MCS}` | | :math:`Q_m`      | | :math:`R\times[1024]` | |                     |
        +===================+====================+=========================+=======================+
        | 0                 | 2                  | 30                      | 0.0586                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 1                 | 2                  | 40                      | 0.0781                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 2                 | 2                  | 50                      | 0.0977                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 3                 | 2                  | 64                      | 0.1250                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 4                 | 2                  | 78                      | 0.1523                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 5                 | 2                  | 99                      | 0.1934                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 6                 | 2                  | 120                     | 0.2344                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 7                 | 2                  | 157                     | 0.3066                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 8                 | 2                  | 193                     | 0.3770                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 9                 | 2                  | 251                     | 0.4902                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 10                | 2                  | 308                     | 0.6016                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 11                | 2                  | 379                     | 0.7402                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 12                | 2                  | 449                     | 0.8770                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 13                | 2                  | 526                     | 1.0273                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 14                | 2                  | 602                     | 1.1758                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 15                | 4                  | 340                     | 1.3281                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 16                | 4                  | 378                     | 1.4766                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 17                | 4                  | 434                     | 1.6953                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 18                | 4                  | 490                     | 1.9141                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 19                | 4                  | 553                     | 2.1602                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 20                | 4                  | 616                     | 2.4063                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 21                | 6                  | 438                     | 2.5564                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 22                | 6                  | 466                     | 2.7305                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 23                | 6                  | 517                     | 3.0293                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 24                | 6                  | 567                     | 3.3223                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 25                | 6                  | 616                     | 3.6094                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 26                | 6                  | 666                     | 3.9023                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 27                | 6                  | 719                     | 4.2129                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 28                | 6                  | 772                     | 4.5234                |
        +-------------------+--------------------+-------------------------+-----------------------+


    .. table:: MCS Index Table 4 (Table 5.1.3.1-4 in [3GPP38214]_)
        :align: center

        +-------------------+--------------------+-------------------------+-----------------------+
        | | MCS Index       | | Modulation Order | | Target Coderate       | | Spectral Efficiency |
        | | :math:`I_{MCS}` | | :math:`Q_m`      | | :math:`R\times[1024]` | |                     |
        +===================+====================+=========================+=======================+
        | 0                 | 2                  | 120                     | 0.2344                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 1                 | 2                  | 193                     | 0.3770                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 2                 | 2                  | 449                     | 0.8770                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 3                 | 4                  | 378                     | 1.4766                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 4                 | 4                  | 490                     | 1.9141                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 5                 | 4                  | 616                     | 2.4063                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 6                 | 6                  | 466                     | 2.7305                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 7                 | 6                  | 517                     | 3.0293                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 8                 | 6                  | 567                     | 3.3223                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 9                 | 6                  | 616                     | 3.6094                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 10                | 6                  | 666                     | 3.9023                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 11                | 6                  | 719                     | 4.2129                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 12                | 6                  | 772                     | 4.5234                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 13                | 6                  | 822                     | 4.8154                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 14                | 6                  | 873                     | 5.1152                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 15                | 8                  | 682.5                   | 5.3320                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 16                | 8                  | 711                     | 5.5547                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 17                | 8                  | 754                     | 5.8906                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 18                | 8                  | 797                     | 6.2266                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 19                | 8                  | 841                     | 6.5703                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 20                | 8                  | 885                     | 6.9141                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 21                | 8                  | 916.5                   | 7.1602                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 22                | 8                  | 948                     | 7.4063                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 23                | 10                 | 805.5                   | 7.8662                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 24                | 10                 | 853                     | 8.3301                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 25                | 10                 | 900.5                   | 8.7939                |
        +-------------------+--------------------+-------------------------+-----------------------+
        | 26                | 10                 | 948                     | 9.2578                |
        +-------------------+--------------------+-------------------------+-----------------------+

    """

    def __init__(self, **kwargs):
        self._name = "Transport Block Configuration"
        super().__init__(**kwargs)
        self.check_config()

    #-----------------------------#
    #---Configurable parameters---#
    #-----------------------------#

    @property
    def mcs_index(self):
        r"""Modulation and coding scheme (MCS) index (denoted as :math:`I_{MCS}`
        in [3GPP38214]_)"""
        self._ifndef("mcs_index", 14) # 16-QAM, r=0.54
        return self._mcs_index

    @mcs_index.setter
    def mcs_index(self, value):
        assert value in range(29), \
            "mcs_index must be in range from 0 to 28."
        self._mcs_index = value

    @property
    def mcs_table(self):
        r"""Indicates which MCS table from [3GPP38214]_ to use. Starts with "1".
        """
        self._ifndef("mcs_table", 1)
        return self._mcs_table

    @mcs_table.setter
    def mcs_table(self, value):
        assert value in range(1,5), \
            "mcs_table must be in range from 1 to 4"
        self._mcs_table = value

    @property
    def channel_type(self):
        r"""5G NR physical channel type. Valid choices are "PDSCH" and "PUSCH".
        """
        self._ifndef("channel_type", "PUSCH")
        return self._channel_type

    @channel_type.setter
    def channel_type(self, value):
        assert value in ("PUSCH", "PDSCH"), \
            'Only "PUSCH" and "PDSCH are supported'
        self._channel_type = value

    @property
    def n_id(self):
        r"""
        `int`, None (default), [0, 1023] : Data scrambling initialization
            :math:`n_\text{ID}`. Data Scrambling ID related to cell id and
            provided by higher layer. If `None`, the
            :class:`~sionna.phy.nr.PUSCHConfig` will automatically set
            :math:`n_\text{ID}=N_\text{ID}^{cell}`.
        """
        self._ifndef("n_id", None)
        return self._n_id

    @n_id.setter
    def n_id(self, value):
        if value is None:
            self._n_id = None
        else:
            assert value in range(1024), \
                "n_id must be in range from 0 to 1023"
            self._n_id = value

    ###
    ### Derived (read-only) parameters
    ###

    @property
    def name(self):
        return "Transport Block Configuration"

    @property
    def target_coderate(self):
        r"""
        `float`, read-only: Target coderate of the TB as defined by the selected
        MCS"""
        _, r = decode_mcs_index(self._mcs_index,
                                self._mcs_table,
                                is_pusch=self._channel_type=='PUSCH')
        return r.numpy()

    @property
    def num_bits_per_symbol(self):
        r"""
        `int`, read-only: Modulation order as defined by the selected MCS"""
        m, _ = decode_mcs_index(self._mcs_index,
                                self._mcs_table,
                                is_pusch=self._channel_type=='PUSCH')
        return m.numpy()

    @property
    def tb_scaling(self):
        r"""
        `float`, 1. (default), read-only: TB scaling factor for PDSCH as
            defined in [3GPP38214]_ Tab. 5.1.3.2-2
        """
        return 1. # only 1. supported at the moment

    #-------------------#
    #---Class methods---#
    #-------------------#

    def check_config(self):
        """Test if configuration is valid"""
        attr_list = ["mcs_index",
                     "mcs_table",
                     "channel_type",
                     "n_id"
                    ]
        for attr in attr_list:
            value = getattr(self, attr)
            setattr(self, attr, value)
